# -*-
import numpy as np
import scipy as sc

def Solve_Linear_Static(K,K_,F_):
    """
    求解KU=F的静力学线性问题,解出U,F_rec;
    
    考虑到求解性能,将使用矩阵规模小于1000使用python计算,否则使用使用eigen编译的pyd计算;

    args:
    
    K(np.ndarray): 整体刚度矩阵,shape=(n,n),奇异矩阵
    K_(np.ndarray): 全局刚度矩阵(经过边界条件处理了的)
    F_(np.ndarray): 全局外力向量,shape=(n,1) or (n,)

    return:
    
    U: 全局位移向量.如果F_为(n,1),则U也是(n,1),反之F_为(n,),则U也是(n,)
    """
    N=K.shape[0]
    if K_.shape[0]!=K_.shape[1]:
        raise ValueError("K必须是方阵")
    if K_.shape[0]!=F_.shape[0]:
        raise ValueError("K和F的行数必须相同")
    U=np.linalg.solve(K_,F_)
    F=K@U
    return U,F 

def AddBcsByReorderK(K,bcs,forces):
    """
    使用划行划列法计算加入边界条件后的整体刚度矩阵K',整体节点力向量F'
    
    args:
        K (numpy.ndarray): 整体刚度矩阵,shape=(n,n),奇异矩阵
        bcs (numpy.ndarray): 边界条件数组,shape=(m,2),每一行代表一个边界条件,第1列代表自由度索引(0-based),第2列代表边界条件值
        forces (numpy.ndarray): 整体节点力向量,shape=(p,2),每一行代表一个节点力,第1列代表自由度索引(0-based),第2列代表节点力值
    return:
        K_ (numpy.ndarray): 加入边界条件后的整体刚度矩阵,shape=(n,n),非奇异矩阵
        F_ (numpy.ndarray): 加入边界条件后的整体节点力向量,shape=(n,1)
    """
    # 初始化向量
    K_=K.copy()
    N=K.shape[0]
    F_=np.zeros(N)
    
    # 已知节点力加到 F_
    for i in range(forces.shape[0]):
        F_[int(forces[i,0])]+=forces[i,1]
    # 循环所有的bcs
    for n in range(bcs.shape[0]):
        j,Uval=int(bcs[n,0]),bcs[n,1]
        #Step1
        for i in range(N):
            F_[i]=F_[i]-K_[i,j]*Uval
        #Step2
        for k in range(N):
            K_[k,j]=0
            K_[j,k]=0
        K_[j,j]=1
        #Step3
        F_[j]=Uval
    return K_,F_

def AddBcsByBigNumber(K,bcs,forces):
    """
    使用充大数(1e8)法计算加入边界条件后的整体刚度矩阵K_,整体节点力向量F_; U=K_^(-1)F_ && KU=F
    
    args:
        K (numpy.ndarray): 整体刚度矩阵,shape=(n,n),奇异矩阵
        bcs (numpy.ndarray): 边界条件数组,shape=(m,2),每一行代表一个边界条件,第1列代表自由度索引(0-based),第2列代表边界条件值
        forces (numpy.ndarray): 整体节点力向量,shape=(p,2),每一行代表一个节点力,第1列代表自由度索引(0-based),第2列代表节点力值
    
    return:
        K_ (numpy.ndarray): 加入边界条件后的整体刚度矩阵,shape=(n,n),非奇异矩阵
        
        F_ (numpy.ndarray): 加入边界条件后的整体节点力向量,shape=(n,1)
    
    优缺点:该方法通用性强,适合大多数的静力学线性问题,但数值精度与大数的取值有关,太小了精度差,太大了容易出现"矩阵奇异"的现象
    """
    # 初始化变量
    C=np.max(K)*1e10
    K_=K.copy()
    N=K.shape[0]
    F_=np.zeros(N)
    # 已知节点力加到 F_
    for i in range(forces.shape[0]):
        F_[int(forces[i,0])]+=forces[i,1]
    # 循环所有位移约束
    for i in range(bcs.shape[0]):
        j,Uval=int(bcs[i,0]),bcs[i,1]
        # Step1
        K_[j,j]=K_[j,j]*C
        # Step2
        F_[j]=K_[j,j]*Uval
    # 求解 K_U=F_
    return K_,F_

def solve_GEigs(K,M,N):
    """
    计算广义特征值问题K*x=lamda*M*x的前N阶特征值&特征向量;返回的特征向量已经质量归一,满足：vectors.T@M@vectors=I
    
    return:
        vals(N,) np.ndarray:特征值
        vectors(n,N) np.ndarray :特征向量
    """
    from scipy.linalg import eigh
    vals,vectors=eigh(a=K,b=M,type=1,subset_by_index=[0,N-1])
    return vals,vectors

def Solve_Natrual_Frequency(K,M,bcs,Nev):
    """
    计算结构的固有频率和模态特征向量
    
    args:
        K (np.ndarray): 全局刚度矩阵,shape=(n,n),奇异矩阵
        M (np.ndarray): 全局质量矩阵,shape=(n,n),可以是一致质量矩阵或集中质量矩阵
        bcs (np.ndarray): 边界条件数组,shape=(m,), 约束自由度索引
        N (int): 需要的前N阶模态
    return:
        EigValues (np.ndarray): 模态特征频率(Hz), shape=(Nev,)
        EigVector (np.ndarray): 模态特征向量,shape=(n,Nev),采用质量归一化
    """
    bcs_=bcs.copy()
    # 提取K，M的子矩阵
    free_dof_ind=np.array(list(set(range(K.shape[0]))-set(bcs_.tolist())),dtype=int)
    Kff=K[np.ix_(free_dof_ind,free_dof_ind)].copy()
    Mff=M[np.ix_(free_dof_ind,free_dof_ind)].copy()
    
    # 求解广义特征值问题Kff*x=lamda*Mff*x
    EigValues,vecs=solve_GEigs(Kff,Mff,Nev)
    
    # 全局的特征值频率
    EigValues=np.sqrt(EigValues)/(2*np.pi)
    
    # 全局结构振型
    EigVectors=np.zeros((K.shape[0],Nev))
    for i in range(Nev):  # 遍历1~Nev模态
        EigVectors[free_dof_ind,i]=vecs[:,i]
    
    return EigValues,EigVectors
    









